#include <vector>
#include <stdint.h>
#include <string>
#include <unordered_map>
#include <mutex>
#include <sstream>
#include "ATen/optional.h"
#include "torch/csrc/assertions.h"
#include "torch/csrc/jit/interned_strings.h"
#include "string.h"
#include <iostream>

namespace torch { namespace jit {

struct InternedStrings {
  InternedStrings()
  : sym_to_info_(static_cast<size_t>(_keys::num_symbols)) {
    #define REGISTER_SYMBOL(n, s) \
      string_to_sym_[#n "::" #s] = n::s; \
      sym_to_info_[n::s] = {namespaces::n, #n "::" #s, #s};

    FORALL_NS_SYMBOLS(REGISTER_SYMBOL)
    #undef REGISTER_SYMBOL
  }
  Symbol symbol(const std::string & s) {
    std::lock_guard<std::mutex> guard(mutex_);
    return _symbol(s);
  }
  std::pair<const char *, const char *> string(Symbol sym) {
    // Builtin Symbols are also in the maps, but
    // we can bypass the need to acquire a lock
    // to read the map for Builtins because we already
    // know their string value
    switch(sym) {
      #define DEFINE_CASE(ns, s) \
        case ns::s: return {#ns "::" #s, #s};
      FORALL_NS_SYMBOLS(DEFINE_CASE)
      #undef DEFINE_CASE
        default:
          return customString(sym);
    }
  }
  Symbol ns(Symbol sym) {
    switch(sym) {
      #define DEFINE_CASE(ns, s) \
        case ns::s: return namespaces::ns;
      FORALL_NS_SYMBOLS(DEFINE_CASE)
      #undef DEFINE_CASE
        default: {
          std::lock_guard<std::mutex> guard(mutex_);
          return sym_to_info_.at(sym).ns;
        }
    }
  }
private:
  // prereq - holding mutex_
  Symbol _symbol(const std::string & s) {
    auto it = string_to_sym_.find(s);
    if(it != string_to_sym_.end())
      return it->second;

    auto pos = s.find("::");
    if(pos == std::string::npos) {
      throw std::runtime_error("all symbols must have a namespace, <namespace>::<string>");
    }
    Symbol ns = _symbol("namespaces::" + s.substr(0, pos));

    Symbol sym(sym_to_info_.size());
    string_to_sym_[s] = sym;
    sym_to_info_.push_back({ns, s, s.substr(pos + strlen("::"))});
    return sym;
  }

  std::pair<const char *, const char *> customString(Symbol sym) {
    std::lock_guard<std::mutex> guard(mutex_);
    SymbolInfo& s = sym_to_info_.at(sym);
    return {s.qual_name.c_str(), s.unqual_name.c_str()};
  }
  std::unordered_map<std::string, Symbol> string_to_sym_;

  struct SymbolInfo {
    Symbol ns;
    std::string qual_name;
    std::string unqual_name;
  };
  std::vector<SymbolInfo> sym_to_info_;

  std::mutex mutex_;
};

static InternedStrings & globalStrings() {
  static InternedStrings s;
  return s;
}

Symbol Symbol::fromQualString(const std::string & s) {
  return globalStrings().symbol(s);
}

const char * Symbol::toUnqualString() const {
  return globalStrings().string(*this).second;
}

const char * Symbol::toQualString() const {
  return globalStrings().string(*this).first;
}

const char * Symbol::toDisplayString() const {
  // TODO: Make this actually return something that's "user friendly".
  // The trouble is that, for this to be usable in printf-style assert
  // statements, this has to return a const char* (whose lifetime is
  // global), so we can't actually assemble a string on the fly.
  return toQualString();
}

Symbol Symbol::ns() const {
  return globalStrings().ns(*this);
}

std::string Symbol::domainString() const {
  return domain_prefix + ns().toUnqualString();
}

Symbol Symbol::fromDomainAndUnqualString(const std::string & d, const std::string & s) {
  if (d.compare(0, domain_prefix.size(), domain_prefix) != 0) {
    std::ostringstream ss;
    ss << "Symbol: domain string is expected to be prefixed with '"
       << domain_prefix << "', e.g. 'org.pytorch.aten'";
    throw std::runtime_error(ss.str());
  }
  std::string qualString = d.substr(domain_prefix.size()) + "::" + s;
  return fromQualString(qualString);
}

}}
